# -*- coding: utf-8 -*-
"""
Created on Wed Aug  3 19:38:01 2022

@author: Shang'xiang
"""

import ttcrpy.rgrid as rg
import numpy as np
import matplotlib.pyplot as plt
# import time

# 创建网格
x = np.arange(0,11.0)
z = np.arange(0,12.0)

# 创建速度模型
v = 2000*np.ones((x.size-1,z.size-1))
v1 = 4000*np.ones((4, 4))

v[2:6,2:6] = v1

fig, ax = plt.subplots()

cs = plt.pcolor(v,cmap='jet')
fig.colorbar(cs)
plt.show()

# 给定发射点和接收点的坐标
# srcs = np.array([[1,1.5],
#                 [1,2.5],
#                 [1,3.5],
#                 [1,4.5],
#                 [1,5.5],
#                 [1,6.5],
#                 [1,7.5],
#                 [1,8.5]])
src = np.array([[1,1.5]])
rcv = np.array([[10,1.5],
                [10,2.5],
                [10,3.5],
                [10,4.5],
                [10,5.5],
                [10,6.5],
                [10,7.5],
                [10,8.5]])

 # 离散网格
grid = rg.Grid2d(x, z, cell_slowness=True, method = 'SPM')

# 速度转换为慢度
slowness = 1./v

tt_all = np.empty(0)
rays_all = list()
# for s in srcs:
#     src = np.array([s])
#     tt, rays = grid.raytrace(src, rcv, slowness, return_rays=True)
#     tt_all = np.append(tt_all,tt)
#     rays_all = rays_all + rays


tt, rays, L = grid.raytrace(src, rcv, slowness, compute_L=True, return_rays=True)
# 绘制走时图
plt.figure(2)
plt.plot(tt_all, 'r-o')
plt.show()

# 保存走时
np.savetxt('G_FSM.txt', tt_all)

plt.figure(3)
# 绘制射线路径图
for r in rays_all:
    plt.plot(r[:,0],r[:,1],'r-s')
    # plt.show()
    # plt.hold(true)
ax = plt.gca()
miloc = plt.MultipleLocator(1)
ax.xaxis.set_minor_locator(miloc)
plt.grid(linestyle = ':', color = 'black',which = 'both')
plt.xlim(0,10)
plt.ylim(0,11)
plt.show()
